#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import os
import PIL
import contextlib
import pytesseract
from .._util import Util, VirtualUser
from .._errors import BuildError


class UtilXp:

    @contextlib.contextmanager
    @staticmethod
    def mountImageFile(workDirObj):
        ret = workDirObj.mount_image_file()
        try:
            yield ret
        finally:
            workDirObj.umount_image_file()


class ScreenshotRecognizerXp:

    def __init__(self):
        self._imgLeftBottom = None

        self._imgTrayIconNewHardware = None

        self._imgWindowTemplateLeftTopCorner = None
        self._imgWindowTemplateRightTop = None
        self._imgWindowTemplateLeftBottomCorner = None
        self._imgWindowTemplateRightBottomCorner = None
        self._imgWindowTemplateLeftEdgeList = None
        self._imgWindowTemplateRightEdgeList = None
        self._imgWindowTemplateBottomEdgeList = None

    def hasStartMenu(self, img):
        if self._imgLeftBottom is None:
            selfDir = os.path.dirname(os.path.realpath(__file__))
            self._imgLeftBottom = PIL.Image.open(os.path.join(selfDir, "final-left-bottom.png"))
        return Util.imgMatchSubImg(img, self._imgLeftBottom, "left-bottom")

    def hasNewHardwareTrayIcon(self, img):
        if self._imgTrayIconNewHardware is None:
            selfDir = os.path.dirname(os.path.realpath(__file__))
            self._imgTrayIconNewHardware = PIL.Image.open(os.path.join(selfDir, "trayicon-new-hardware.png"))

        region = (
            self._imgLeftBottom.size[1],
            img.size[1] - 7 - self._imgTrayIconNewHardware.size[1],
            img.size[0],
            img.size[1] - 7
        )
        return Util.imgSearchSubImg(img, self._imgTrayIconNewHardware, region=region) is not None

    def getWindowRect(self, img, includeDecoration=True):
        if self._imgWindowTemplateLeftTopCorner is None:
            selfDir = os.path.dirname(os.path.realpath(__file__))
            imgTmpl = PIL.Image.open(os.path.join(selfDir, "template-window.png"))
            maskColor = (0, 255, 0)

            self._imgWindowTemplateLeftTopCorner = Util.imgCropRectByExcludingPureColor(imgTmpl, 0, 0, maskColor)
            self._imgWindowTemplateRightTopCorner = Util.imgCropRectByExcludingPureColor(imgTmpl, imgTmpl.width - 1, 0, maskColor)
            self._imgWindowTemplateLeftBottomCorner = Util.imgCropRectByExcludingPureColor(imgTmpl, 0, imgTmpl.height - 1, maskColor)
            self._imgWindowTemplateRightBottomCorner = Util.imgCropRectByExcludingPureColor(imgTmpl, imgTmpl.width - 1, imgTmpl.height - 1, maskColor)

            self._imgWindowTemplateLeftEdgeList = []
            if True:
                y = self._imgWindowTemplateLeftTopCorner.height + 1     # extra "+ 1" is to jump over the first padding
                while True:
                    tImg = Util.imgCropRectByExcludingPureColor(imgTmpl, 0, y, maskColor, bStrict=False)
                    if tImg is None:
                        break
                    assert all([tImg.size == t.size and tImg != t for t in self._imgWindowTemplateLeftEdgeList])
                    self._imgWindowTemplateLeftEdgeList.append(tImg)
                    y += tImg.height + 1                                # extra "+ 1" is to jump over padding
                assert y < imgTmpl.height - self._imgWindowTemplateLeftBottomCorner.height

            self._imgWindowTemplateRightEdgeList = []
            if True:
                y = self._imgWindowTemplateRightTopCorner.height + 1    # extra "+ 1" is to jump over the first padding
                while True:
                    tImg = Util.imgCropRectByExcludingPureColor(imgTmpl, imgTmpl.width - 1, y, maskColor, bStrict=False)
                    if tImg is None:
                        break
                    assert all([tImg.size == t.size and tImg != t for t in self._imgWindowTemplateRightEdgeList])
                    self._imgWindowTemplateRightEdgeList.append(tImg)
                    y += tImg.height + 1                                # extra "+ 1" is to jump over padding
                assert y < imgTmpl.height - self._imgWindowTemplateRightBottomCorner.height

            self._imgWindowTemplateBottomEdgeList = []
            if True:
                x = self._imgWindowTemplateLeftBottomCorner.width + 1   # extra "+ 1" is to jump over the first padding
                while True:
                    tImg = Util.imgCropRectByExcludingPureColor(imgTmpl, x, imgTmpl.height - 1, maskColor, bStrict=False)
                    if tImg is None:
                        break
                    assert all([tImg.size == t.size and tImg != t for t in self._imgWindowTemplateBottomEdgeList])
                    self._imgWindowTemplateBottomEdgeList.append(tImg)
                    x += tImg.width + 1                                 # extra "+ 1" is to jump over padding
                assert x < imgTmpl.width - self._imgWindowTemplateRightBottomCorner.width

        left = None
        top = None
        right = None
        bottom = None

        if img.width < self._imgWindowTemplateRightTopCorner.width or img.height < self._imgWindowTemplateRightTopCorner.height:
            return None

        ret = Util.imgSearchSubImg(img, self._imgWindowTemplateRightTopCorner, search_from="right-top")
        if ret is None:
            return None
        right = ret[0] + self._imgWindowTemplateRightTopCorner.width
        top = ret[1]

        ret = self._findAlongTheEdge(img, self._imgWindowTemplateRightTopCorner, right, top, 'right', 'down', self._imgWindowTemplateRightEdgeList, self._imgWindowTemplateRightBottomCorner)
        if ret is None:
            return None
        assert right == ret[0]
        bottom = ret[1]

        ret = self._findAlongTheEdge(img, self._imgWindowTemplateRightBottomCorner, right, bottom, 'bottom', 'left', self._imgWindowTemplateBottomEdgeList, self._imgWindowTemplateLeftBottomCorner)
        if ret is None:
            return None
        left = ret[0]
        assert bottom == ret[1]

        ret = self._findAlongTheEdge(img, self._imgWindowTemplateLeftBottomCorner, left, bottom, 'left', 'up', self._imgWindowTemplateLeftEdgeList, self._imgWindowTemplateLeftTopCorner)
        if ret is None:
            return None
        assert left == ret[0]
        assert top == ret[1]

        if includeDecoration:
            return (left, top, right, bottom)
        else:
            return (left + self._imgWindowTemplateLeftEdgeList[0].width, top + self._imgWindowTemplateLeftTopCorner.height, right - self._imgWindowTemplateRightEdgeList[0].width, bottom - self._imgWindowTemplateBottomEdgeList[0].height)

    @staticmethod
    def _findAlongTheEdge(img, imgStartCorner, sx, sy, side, direction, imgEdgeList, imgEndCorner):
        if side == "left":
            if direction == "down":
                x = sx                          # sx is "left"
                y = sy + imgStartCorner.height  # sy is "top"
                while y <= img.height - imgEndCorner.height:
                    if Util.imgSearchSubImg(img, imgEndCorner, region=(x, y, x + imgEndCorner.width, y + imgEndCorner.height)) is not None:
                        return (x, y + imgEndCorner.height)         # left-bottom corner found, even if no edge exists, returns (left, bottom)
                    y += 1
                    # bFound = False
                    # for imgEdge in imgEdgeList:
                    #     if Util.imgSearchSubImg(img, imgEdge, region=(x, y, x + imgEdge.width, y + imgEdge.height)) is not None:
                    #         y += imgEdge.height
                    #         bFound = True
                    #         break
                    # if not bFound:
                    #     return None                                 # edge not match
                return None                                         # left-bottom corner not found
            if direction == "up":
                x = sx                          # sx is "left"
                y = sy - imgStartCorner.height  # sy is "bottom"
                while y >= imgEndCorner.height:
                    if Util.imgSearchSubImg(img, imgEndCorner, region=(x, y - imgEndCorner.height, x + imgEndCorner.width, y)) is not None:
                        return (x, y - imgEndCorner.height)         # left-top corner found, even if no edge exists, returns (left, top)
                    y -= 1
                    # bFound = False
                    # for imgEdge in imgEdgeList:
                    #     if Util.imgSearchSubImg(img, imgEdge, region=(x, y - imgEdge.height, x + imgEdge.width, y)) is not None:
                    #         y -= imgEdge.height
                    #         bFound = True
                    #         break
                    # if not bFound:
                    #     print("ddd1")
                    #     return None                                 # edge not match
                return None                                         # left-top corner not found
            assert False

        if side == "right":
            if direction == "down":
                x = sx                          # sx is "right"
                y = sy + imgStartCorner.height  # sy is "top"
                while y <= img.height - imgEndCorner.height:
                    if Util.imgSearchSubImg(img, imgEndCorner, region=(x - imgEndCorner.width, y, x, y + imgEndCorner.height)) is not None:
                        return (x, y + imgEndCorner.height)         # right-bottom corner found, even if no edge exists, returns (right, bottom)
                    y += 1
                    # bFound = False
                    # for imgEdge in imgEdgeList:
                    #     if Util.imgSearchSubImg(img, imgEdge, region=(x - imgEdge.width, y, x, y + imgEdge.height)) is not None:
                    #         y += imgEdge.height
                    #         bFound = True
                    #         break
                    # if not bFound:
                    #     print("ddd2")
                    #     return None                                 # edge not match
                return None                                         # right-bottom corner not found
            if direction == "up":
                x = sx                          # sx is "right"
                y = sy - imgStartCorner.height  # sy is "bottom"
                while y >= imgEndCorner.height:
                    if Util.imgSearchSubImg(img, imgEndCorner, region=(x - imgEndCorner.width, y - imgEndCorner.height, x, y)) is not None:
                        return (x, y - imgEndCorner.height)         # right-top corner found, even if no edge exists, returns (right, top)
                    y -= 1
                    # bFound = False
                    # for imgEdge in imgEdgeList:
                    #     if Util.imgSearchSubImg(img, imgEdge, region=(x - imgEdge.width, y - imgEdge.height, x, y)) is not None:
                    #         y -= imgEdge.height
                    #         bFound = True
                    #         break
                    # if not bFound:
                    #     return None                                 # edge not match
                return None                                         # right-top corner not found
            assert False

        if side == "top":
            # top side is not findable because window title bar has dynamic text content and icon
            assert False

        if side == "bottom":
            if direction == "right":
                x = sx + imgStartCorner.width   # sx is "left"
                y = sy                          # sy is "bottom"
                while x <= img.width - imgEndCorner.width:
                    if Util.imgSearchSubImg(img, imgEndCorner, region=(x, y - imgEndCorner.height, x + imgEndCorner.width, y)) is not None:
                        return (x + imgEndCorner.width, y)          # right-bottom corner found, even if no edge exists, returns (right, bottom)
                    x += 1
                    # bFound = False
                    # for imgEdge in imgEdgeList:
                    #     if Util.imgSearchSubImg(img, imgEdge, region=(x, y - imgEdge.height, x + imgEdge.width, y)) is not None:
                    #         x += imgEdge.width
                    #         bFound = True
                    #         break
                    # if not bFound:
                    #     return None                                 # edge not match
                return None                                         # right-bottom corner not found
            if direction == "left":
                x = sx - imgStartCorner.width   # sx is "right"
                y = sy                          # sy is "bottom"
                while x >= imgEndCorner.width:
                    if Util.imgSearchSubImg(img, imgEndCorner, region=(x - imgEndCorner.width, y - imgEndCorner.height, x, y)) is not None:
                        return (x - imgEndCorner.width, y)          # left-bottom corner found, even if no edge exists, returns (left, bottom)
                    x -= 1
                    # bFound = False
                    # for imgEdge in imgEdgeList:
                    #     if Util.imgSearchSubImg(img, imgEdge, region=(x - imgEdge.width, y - imgEdge.height, x, y)) is not None:
                    #         x -= imgEdge.width
                    #         bFound = True
                    #         break
                    # if not bFound:
                    #     print("ddd3")
                    #     return None                                 # edge not match
                return None                                         # left-bottom corner not found

        assert False


class VirtualUserXp(VirtualUser):

    class BootFromCdrom:

        def do(self, parent, img):
            # OCR result should be:
            #
            # Booting from DVD/CD...
            #
            # Press any key to boot from CD.
            #

            s = pytesseract.image_to_string(img, lang=parent._ocrLang)
            if "Press any key to boot from CD" in s:
                parent._vm.input_keyboard("keys", "kp_enter")
                return True
            return False

    class WaitFirstBoot:

        def __init__(self):
            self._lastImg = None
            self._stableCount = None

        def do(self, parent, img):
            if True:
                # OCR result should be:
                #
                # System Settings Change
                #
                # Windows has finished installing new devices.
                #
                # XXXXXX
                #
                # Do you want to restart your computer now?
                #

                s = pytesseract.image_to_string(img, lang=parent._ocrLang)
                if "Do you want to restart your computer now?" in s:
                    # press enter to reboot
                    parent._vm.input_mouse("abs-move-and-click-left-button", img.size[0] // 2, img.size[1] // 2)    # click at the center of screen, focus "System Settings Change" dialog
                    parent._vm.input_keyboard("keys", "kp_enter")
                    return False

            if self._lastImg is None:
                if not parent._scr.hasStartMenu(img):
                    # start menu is not there
                    return False
                if parent._scr.hasNewHardwareTrayIcon(img):
                    # new hardware icon is in the tray
                    return False
                self._lastImg = img
                self._stableCount = 0
                return False

            if self._lastImg != img:
                self._lastImg = None
                self._stableCount = None
                return False

            self._stableCount += 1
            if self._stableCount > 10:
                # the system has stablized for 10 seconds
                return True

            return False

    def __init__(self, ocr_lang, vm):
        super().__init__(vm, BuildError)
        self._ocrLang = ocr_lang
        self._scr = ScreenshotRecognizerXp()
